package com.kunge2013.chatbi.embedding;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.DimensionAwareEmbeddingModel;
import dev.langchain4j.model.output.Response;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.AbstractEmbeddingModel;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.List;


/**
 * @Author kunge2013
 * @Description 自定义EmbeddingModel
 * @Date 2024/11/10 13:36
 * @Version 1.0
 */
public class SelfEmbeddingModel extends AbstractEmbeddingModel {
    private final DimensionAwareEmbeddingModel dimensionAwareEmbeddingModel;
    private final MetadataMode metadataMode;

    public SelfEmbeddingModel(DimensionAwareEmbeddingModel dimensionAwareEmbeddingModel) {
        this.dimensionAwareEmbeddingModel = dimensionAwareEmbeddingModel;
        this.metadataMode = MetadataMode.EMBED;
    }


    @Override
    public EmbeddingResponse call(EmbeddingRequest request) {
        List<org.springframework.ai.embedding.Embedding> embeddings = new ArrayList<>();
        List<String> instructions = request.getInstructions();
        for (int i = 0; i < instructions.size(); i++) {
            String text = instructions.get(i);
            Response<Embedding> response = dimensionAwareEmbeddingModel.embed(text);
            Embedding content = response.content();
            float[] embedding = content.vector();
            var resultEmbd = new org.springframework.ai.embedding.Embedding(embedding, i);
            embeddings.add(resultEmbd);
        }
        return new EmbeddingResponse(embeddings);
    }

    @Override
    public float[] embed(Document document) {
        Assert.notNull(document, "Document must not be null");
        return this.embed(document.getFormattedContent(this.metadataMode));
    }

    @Override
    public int dimensions() {
        return 1024;
    }
}
